from loguru import logger

class TaskStatusManager:
    def __init__(self, root_agent_instance):
        self.root_agent = root_agent_instance

    def pause_task(self, task_id: str):
        if task_id in self.root_agent.task_states:
            self.root_agent.task_states[task_id]["status"] = "paused"
            logger.info(f"Task {task_id} paused.")
        else:
            logger.warning(f"Task {task_id} not found.")

    def resume_task(self, task_id: str):
        if task_id in self.root_agent.task_states:
            self.root_agent.task_states[task_id]["status"] = "running"
            logger.info(f"Task {task_id} resumed.")
        else:
            logger.warning(f"Task {task_id} not found.")

    def cancel_task(self, task_id: str):
        if task_id in self.root_agent.task_states:
            self.root_agent.task_states[task_id]["status"] = "cancelled"
            logger.info(f"Task {task_id} cancelled.")
        else:
            logger.warning(f"Task {task_id} not found.")

    def get_task_status(self, task_id: str):
        return self.root_agent.task_states.get(task_id, {"status": "not_found", "progress": 0, "result": "", "error": ""})
